# app.py

import os
import json
import time
from flask import Flask, render_template, request, jsonify, send_from_directory
from flask_sqlalchemy import SQLAlchemy
from werkzeug.utils import secure_filename
from PIL import Image
import pytesseract
from googletrans import Translator

# 配置OCR引擎路径（根据实际情况修改）
pytesseract.pytesseract.tesseract_cmd = r'/usr/bin/tesseract'

app = Flask(__name__)
app.config['SECRET_KEY'] = 'your-secret-key'
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///documents.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

# # 配置上传文件夹
# UPLOAD_FOLDER = 'uploads'
# if not os.path.exists(UPLOAD_FOLDER):
#     os.makedirs(UPLOAD_FOLDER)

# app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 配置上传文件夹
UPLOAD_FOLDER = 'uploads'
TRANSLATED_FOLDER = 'translated'  # 添加翻译结果文件夹

for folder in [UPLOAD_FOLDER, TRANSLATED_FOLDER]:
    if not os.path.exists(folder):
        os.makedirs(folder)

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['TRANSLATED_FOLDER'] = TRANSLATED_FOLDER

db = SQLAlchemy(app)

# 文件模型
class File(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    filename = db.Column(db.String(255), nullable=False)
    file_type = db.Column(db.String(50), nullable=False)  # pdf, image
    status = db.Column(db.String(50), default='uploaded')  # uploaded, processing, completed, failed
    progress = db.Column(db.Integer, default=0)
    translation_service = db.Column(db.String(50))
    created_at = db.Column(db.DateTime, default=db.func.current_timestamp())
    completed_at = db.Column(db.DateTime)
    original_text = db.Column(db.Text)
    translated_text = db.Column(db.Text)
    ocr_result = db.Column(db.Text)  # 图片OCR结果
    translated_filename = db.Column(db.String(255))  # 翻译结果文件名
    
# 创建数据库表
with app.app_context():
    db.create_all()

# 首页
@app.route('/')
def index():
    return render_template('index.html')

# 获取文件列表
@app.route('/api/files')
def get_files():
    files = File.query.all()
    file_list = []
    
    for file in files:
        file_data = {
            'id': file.id,
            'filename': file.filename,
            'file_type': file.file_type,
            'status': file.status,
            'progress': file.progress,
            'translation_service': file.translation_service,
            'created_at': file.created_at.strftime('%Y-%m-%d %H:%M:%S') if file.created_at else '',
            'completed_at': file.completed_at.strftime('%Y-%m-%d %H:%M:%S') if file.completed_at else ''
        }
        file_list.append(file_data)
    
    return jsonify(file_list)

# 上传并开始翻译文件
@app.route('/upload_and_translate', methods=['POST'])
def upload_and_translate():
    try:
        if 'file' not in request.files:
            return jsonify({'error': '没有文件部分'}), 400
        
        file = request.files['file']
        
        if file.filename == '':
            return jsonify({'error': '没有选择文件'}), 400
        
        # 检查文件类型
        file_ext = os.path.splitext(file.filename)[1].lower()
        if file_ext in ['.pdf']:
            file_type = 'pdf'
        elif file_ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
            file_type = 'image'
        else:
            return jsonify({'error': '不支持的文件类型'}), 400
        
        filename = secure_filename(file.filename)
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)
        
        # 获取翻译服务和API密钥
        service_provider = request.form.get('service_provider', 'doubao')
        api_key = request.form.get('api_key', '')
        
        # 创建文件记录
        new_file = File(
            filename=filename,
            file_type=file_type,
            status='processing',  # 直接设为处理中
            progress=0,
            translation_service=service_provider
        )
        db.session.add(new_file)
        db.session.commit()
        
        # 异步处理翻译任务
        if file_type == 'pdf':
            # 处理PDF文件
            for i in range(101):
                new_file.progress = i
                db.session.commit()
                time.sleep(0.1)  # 模拟处理时间
            
            new_file.status = 'completed'
            new_file.completed_at = db.func.current_timestamp()
            db.session.commit()
            
        elif file_type == 'image':
            # 处理图片文件
            try:
                img = Image.open(file_path)
                ocr_text = pytesseract.image_to_string(img, lang='eng+chi_sim')  # 中英文识别
                new_file.ocr_result = ocr_text
                
                # 翻译OCR结果
                translator = Translator()
                try:
                    translation = translator.translate(ocr_text, dest='zh-cn')
                    new_file.translated_text = translation.text
                except Exception as e:
                    new_file.translated_text = f"翻译失败: {str(e)}"
                
                # 更新进度和状态
                for i in range(101):
                    new_file.progress = i
                    db.session.commit()
                    time.sleep(0.1)  # 模拟处理时间
                
                new_file.status = 'completed'
                new_file.completed_at = db.func.current_timestamp()
                db.session.commit()
                
            except Exception as e:
                new_file.status = 'failed'
                db.session.commit()
                return jsonify({'error': f'处理图片失败: {str(e)}'}), 500
        
        return jsonify({
            'message': '文件上传并开始翻译',
            'filename': filename,
            'file_type': file_type,
            'id': new_file.id
        })
    except Exception as e:
        db.session.rollback()
        return jsonify({'error': f'操作失败: {str(e)}'}), 500

# 获取翻译状态
@app.route('/get_translation_status/<int:file_id>')
def get_translation_status(file_id):
    file = File.query.get(file_id)
    if not file:
        return jsonify({'error': '文件不存在'}), 404
    
    return jsonify({
        'status': file.status,
        'progress': file.progress
    })


# 获取翻译结果
@app.route('/get_translation_result/<int:file_id>')
def get_translation_result(file_id):
    file = File.query.get(file_id)
    if not file:
        return jsonify({'error': '文件不存在'}), 404
    
    if file.file_type == 'pdf':
        # 对于PDF文件，返回模拟的翻译结果
        return jsonify([
            {
                'original': 'This is a sample PDF document for translation.',
                'translated': '这是一个用于翻译的示例PDF文档。'
            },
            {
                'original': 'PDF translation involves extracting text from the document and translating it into another language.',
                'translated': 'PDF翻译涉及从文档中提取文本并将其翻译成另一种语言。'
            }
        ])
    elif file.file_type == 'image':
        # 对于图片文件，返回OCR和翻译结果
        if file.ocr_result and file.translated_text:
            return jsonify([
                {
                    'original': file.ocr_result,
                    'translated': file.translated_text
                }
            ])
        else:
            return jsonify([
                {
                    'original': 'No OCR result available',
                    'translated': '没有可用的OCR结果'
                }
            ])

# 下载原始文件
@app.route('/download_original/<int:file_id>')
def download_original(file_id):
    file = File.query.get(file_id)
    if not file:
        return "File not found", 404
    
    return send_from_directory(app.config['UPLOAD_FOLDER'], file.filename)

# 下载翻译后的文件
@app.route('/download_translated/<int:file_id>')
def download_translated(file_id):
    file = File.query.get(file_id)
    if not file:
        return "File not found", 404
    
    if file.file_type == 'pdf':
        # 对于PDF，返回原始文件（实际应用中应该返回翻译后的PDF）
        return send_from_directory(app.config['UPLOAD_FOLDER'], file.filename)
    elif file.file_type == 'image':
        # 对于图片，返回翻译文本
        if file.translated_text:
            return jsonify({
                'filename': file.filename,
                'translated_text': file.translated_text
            })
        else:
            return "Translation not available", 404

@app.route('/api/delete/<file_id>', methods=['DELETE'])
def delete_file(file_id):
    file = File.query.get(file_id)
    if not file:
        return jsonify({"error": "文件不存在"}), 404
    
    try:
        # 删除物理文件
        original_path = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
        if os.path.exists(original_path):
            os.remove(original_path)
        
        # 删除翻译结果文件
        if file.translated_filename:
            translated_path = os.path.join(app.config['TRANSLATED_FOLDER'], file.translated_filename)
            if os.path.exists(translated_path):
                os.remove(translated_path)
        
        # 从数据库中删除记录
        db.session.delete(file)
        db.session.commit()
        
        return jsonify({"message": "文件已成功删除"}), 200
    except Exception as e:
        db.session.rollback()
        return jsonify({"error": f"删除文件失败: {str(e)}"}), 500
if __name__ == '__main__':
    app.run(debug=True,port=5001)